import heapq
import pandas as pd
from collections import defaultdict
import matplotlib.pyplot as plt
import networkx as nx
import matplotlib
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import tkinter as tk
from tkinter import ttk, messagebox
import traceback

# 保留原有的函数
def read_data_from_excel(file_path):
    print("开始从Excel文件读取数据...")
    
    try:
        # 读取第一个工作表(sheet1)中的距离矩阵
        print("正在从Sheet1读取邻接矩阵...")
        distance_df = pd.read_excel(file_path, sheet_name=0, header=0, index_col=0)
        print(f"邻接矩阵读取成功! 形状: {distance_df.shape}")
        
        # 读取第二个工作表(sheet2)中的启发函数值
        print("正在从Sheet2读取启发函数值...")
        heuristic_df = pd.read_excel(file_path, sheet_name=1)  # sheet_name=1
        print(f"启发函数值读取成功! 形状: {heuristic_df.shape}")
        
        return distance_df, heuristic_df
    
    except Exception as e:
        print(f"读取Excel数据失败: {str(e)}")
        raise

def create_romania_map(distance_df):
    graph = defaultdict(list) # 使用defaultdict简化邻接表的创建
    print("开始构建罗马尼亚地图...")
    connections_count = 0
    
    # 根据距离矩阵构建邻接表
    for city1 in distance_df.index:
        for city2 in distance_df.columns:
            # 跳过自环连接（城市到自己的连接）
            if city1 == city2:
                continue
                
            # 获取距离值并确保它是一个标量
            try:
                distance = distance_df.loc[city1, city2]
                # 使用标量比较方法而不是直接使用 != 运算符
                if isinstance(distance, (int, float)):
                    if distance != 1000 and distance > 0:  # 确保距离大于0
                        graph[city1].append((city2, distance))
                        connections_count += 1
                else:
                    # 如果是Series，转换为标量后比较
                    distance_val = distance.iloc[0] if hasattr(distance, 'iloc') else float(distance)
                    if distance_val != 1000 and distance_val > 0:  # 确保距离大于0
                        graph[city1].append((city2, distance_val))
                        connections_count += 1
            except Exception as e:
                print(f"处理城市 {city1} 到 {city2} 的连接时出错: {e}")
    
    print(f"罗马尼亚地图构建完成! 共有 {len(graph)} 个城市, {connections_count} 条连接")
    return graph
    
def build_heuristic_function(heuristic_df):
    heuristic = {}
    print("开始构建启发函数...")
    
    # 读取城市名和启发函数值
    count = 0
    for i in range(len(heuristic_df)):
        try:
            city = heuristic_df.iloc[i, 0]  # 第一列是城市名
            value = heuristic_df.iloc[i, 1]  # 第二列是启发函数值
            if pd.notna(city) and pd.notna(value):
                heuristic[city] = value
                count += 1
        except Exception as e:
            print(f"处理启发函数的第 {i} 行时出错: {e}")
    
    print(f"启发函数构建完成! 共加载了 {count} 个城市的启发值")
    
    # 打印部分启发函数值作为调试信息
    if heuristic:
        print("部分启发函数值示例:")
        items = list(heuristic.items())[:5]
        for city, value in items:
            print(f"  {city}: {value}")
    
    return heuristic

#1. 代价一致搜索
def uniform_cost_search(graph, start, goal):
    # 优先队列中的元素为 (路径代价, 城市, 路径)
    # 根据路径代价进行排序 路径代价是从起点到当前城市的总代价，即邻接矩阵的值
    priority_queue = [(0, start, [start])] # 初始化优先队列
    visited = set() # 已访问的节点集合
    expanded = 0 # 扩展的节点计数
    
    while priority_queue:
        cost, current, path = heapq.heappop(priority_queue)
        expanded += 1
        
        if current == goal:
            return path, cost, expanded
        
        if current in visited:
            continue
        
        visited.add(current)
        
        for neighbor, step_cost in graph[current]:
            if neighbor not in visited:
                new_cost = cost + step_cost
                new_path = path + [neighbor]
                heapq.heappush(priority_queue, (new_cost, neighbor, new_path))
    
    return None, None, expanded  # 没有找到路径

# 2. 贪婪最佳优先搜索
def greedy_search(graph, start, goal, h):
    # h 是一个字典，包含每个城市到目标城市的启发式距离
    # 优先队列中的元素为 (启发式距离, 城市, 路径, 总代价)
    #根据启发式函数h进行排序
    priority_queue = [(h[start], start, [start], 0)]
    visited = set()
    expanded = 0
    
    while priority_queue:
        _, current, path, cost = heapq.heappop(priority_queue)
        expanded += 1
        
        if current == goal:
            return path, cost, expanded
        
        if current in visited:
            continue
        
        visited.add(current)
        
        for neighbor, step_cost in graph[current]:
            if neighbor not in visited:
                new_cost = cost + step_cost
                new_path = path + [neighbor]
                heapq.heappush(priority_queue, (h[neighbor], neighbor, new_path, new_cost))
    
    return None, None, expanded  # 没有找到路径

#· 3. A*算法
def a_star_search(graph, start, goal, h):
    # 优先队列中的元素为 (总评估代价, 城市, 路径, 已走路径代价)
    # f(n) = g(n) + h(n)，g(n)是已走代价，h(n)是启发式估计，根据f(n)进行排序
    priority_queue = [(h[start], start, [start], 0)]
    visited = set()
    expanded = 0
    
    while priority_queue:
        _, current, path, cost_so_far = heapq.heappop(priority_queue)
        expanded += 1
        
        if current == goal:
            return path, cost_so_far, expanded
        
        if current in visited:
            continue
        
        visited.add(current)
        
        for neighbor, step_cost in graph[current]:
            if neighbor not in visited:
                new_cost = cost_so_far + step_cost
                new_path = path + [neighbor]
                # f(n) = g(n) + h(n)，g(n)是已走代价，h(n)是启发式估计
                f = new_cost + h[neighbor]
                heapq.heappush(priority_queue, (f, neighbor, new_path, new_cost))
    
    return None, None, expanded  # 没有找到路径

# 创建GUI应用
class RomaniaMapApp:
    def __init__(self, root):
        self.root = root
        self.root.title("罗马尼亚度假问题可视化")
        self.root.geometry("1000x700")
        
        # 设置字体
        matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文
        matplotlib.rcParams['axes.unicode_minus'] = False    # 用来正常显示负号
        
        # 加载数据
        self.file_path = "罗马尼亚度假问题地图及启发函数值.xlsx"
        self.load_data()
        
        # 创建UI
        self.create_ui()
        
    def load_data(self):
        try:
            # 读取Excel数据
            self.distance_df, self.heuristic_df = read_data_from_excel(self.file_path)
            
            # 构建罗马尼亚地图
            self.romania_map = create_romania_map(self.distance_df)
            
            # 构建启发函数
            self.heuristic = build_heuristic_function(self.heuristic_df)
            
            # 创建networkx图
            self.G = nx.DiGraph()
            
            # 添加节点和边
            for city in self.romania_map:
                self.G.add_node(city)
                for neighbor, distance in self.romania_map[city]:
                    self.G.add_edge(city, neighbor, weight=distance)
                    
            # 使用布局算法计算节点位置
            # 我们使用kamada_kawai_layout因为它试图保持连接节点之间的距离一致
            self.pos = nx.kamada_kawai_layout(self.G)
            
            # 打印节点位置信息
            print("使用自动布局生成城市位置")
            
        except Exception as e:
            messagebox.showerror("加载错误", f"加载数据时出错: {str(e)}")
            traceback.print_exc()
    
    def create_ui(self):
        # 创建主框架
        main_frame = ttk.Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # 创建控制面板
        control_frame = ttk.LabelFrame(main_frame, text="控制面板")
        control_frame.pack(side=tk.LEFT, fill=tk.Y, padx=5, pady=5)
        
        # 起点终点选择
        ttk.Label(control_frame, text="起点:").grid(row=0, column=0, padx=5, pady=5, sticky=tk.W)
        self.start_var = tk.StringVar(value="Arad")
        start_combo = ttk.Combobox(control_frame, textvariable=self.start_var, 
                                   values=sorted(list(self.romania_map.keys())))
        start_combo.grid(row=0, column=1, padx=5, pady=5, sticky=tk.W)
        
        ttk.Label(control_frame, text="终点:").grid(row=1, column=0, padx=5, pady=5, sticky=tk.W)
        self.goal_var = tk.StringVar(value="Bucharest")
        goal_combo = ttk.Combobox(control_frame, textvariable=self.goal_var, 
                                 values=sorted(list(self.romania_map.keys())))
        goal_combo.grid(row=1, column=1, padx=5, pady=5, sticky=tk.W)
        
        # 算法按钮
        ttk.Label(control_frame, text="选择算法:").grid(row=2, column=0, columnspan=2, padx=5, pady=5, sticky=tk.W)
        
        ucs_btn = ttk.Button(control_frame, text="代价一致搜索", 
                            command=lambda: self.run_algorithm("ucs"))
        ucs_btn.grid(row=3, column=0, columnspan=2, padx=5, pady=5, sticky=tk.W+tk.E)
        
        greedy_btn = ttk.Button(control_frame, text="贪婪最佳优先搜索", 
                               command=lambda: self.run_algorithm("greedy"))
        greedy_btn.grid(row=4, column=0, columnspan=2, padx=5, pady=5, sticky=tk.W+tk.E)
        
        astar_btn = ttk.Button(control_frame, text="A*算法", 
                              command=lambda: self.run_algorithm("astar"))
        astar_btn.grid(row=5, column=0, columnspan=2, padx=5, pady=5, sticky=tk.W+tk.E)
        
        # 布局重置按钮
        reset_btn = ttk.Button(control_frame, text="重新布局", 
                              command=self.regenerate_layout)
        reset_btn.grid(row=6, column=0, columnspan=2, padx=5, pady=5, sticky=tk.W+tk.E)
        
        # 结果框
        result_frame = ttk.LabelFrame(control_frame, text="搜索结果")
        result_frame.grid(row=7, column=0, columnspan=2, padx=5, pady=5, sticky=tk.W+tk.E)
        
        self.result_text = tk.Text(result_frame, width=30, height=10, wrap=tk.WORD)
        self.result_text.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 地图显示区域
        map_frame = ttk.LabelFrame(main_frame, text="罗马尼亚地图")
        map_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 创建matplotlib图形
        self.fig, self.ax = plt.subplots(figsize=(8, 7))
        self.canvas = FigureCanvasTkAgg(self.fig, master=map_frame)
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
        
        # 绘制初始地图
        self.draw_map()
    
    def regenerate_layout(self):
        """重新生成图形布局"""
        # 使用不同的布局算法
        layout_algorithms = [
            ("spring_layout", nx.spring_layout),
            ("kamada_kawai_layout", nx.kamada_kawai_layout),
            ("circular_layout", nx.circular_layout),
            ("shell_layout", nx.shell_layout),
            ("spectral_layout", nx.spectral_layout)
        ]
        
        # 循环使用下一个布局算法
        current_layout = getattr(self, "current_layout", -1)
        current_layout = (current_layout + 1) % len(layout_algorithms)
        self.current_layout = current_layout
        
        layout_name, layout_func = layout_algorithms[current_layout]
        self.pos = layout_func(self.G)
        
        messagebox.showinfo("布局更新", f"已切换到 {layout_name} 布局")
        
        # 重绘地图
        self.draw_map()
        
    def draw_map(self, path=None):
        self.ax.clear()
        
        # 创建边列表，排除自环和权重为0的边
        edgelist = [(u, v) for u, v, d in self.G.edges(data=True) if u != v and d['weight'] > 0]
        
        # 绘制边
        edge_colors = []
        edge_widths = []
        
        for u, v in edgelist:
            if path and u in path and v in path and path.index(u) == path.index(v) - 1:
                # 如果边在路径中，设为红色
                edge_colors.append('red')
                edge_widths.append(2.5)
            else:
                edge_colors.append('black')
                edge_widths.append(1.0)
        
        # 绘制边
        nx.draw_networkx_edges(self.G, self.pos, edgelist=edgelist, ax=self.ax, 
                          edge_color=edge_colors, width=edge_widths)
        
        # 绘制节点
        node_colors = []
        node_sizes = []
        
        for node in self.G.nodes():
            if path and node in path:
                if node == path[0]:  # 起点
                    node_colors.append('green')
                    node_sizes.append(700)
                elif node == path[-1]:  # 终点
                    node_colors.append('red')
                    node_sizes.append(700)
                else:  # 路径中间点
                    node_colors.append('orange')
                    node_sizes.append(500)
            else:  # 非路径节点
                node_colors.append('skyblue')
                node_sizes.append(300)
        
        nx.draw_networkx_nodes(self.G, self.pos, ax=self.ax, 
                              node_color=node_colors, node_size=node_sizes)
        
        # 绘制标签
        nx.draw_networkx_labels(self.G, self.pos, ax=self.ax, font_size=9, font_family='SimHei')
        
        # 绘制边标签（距离），排除自环和权重为0的边
        edge_labels = {(u, v): d['weight'] for u, v, d in self.G.edges(data=True) 
                      if u != v and d['weight'] > 0}
        nx.draw_networkx_edge_labels(self.G, self.pos, edge_labels=edge_labels, font_size=8)
        
        # 设置图形参数
        self.ax.set_title("罗马尼亚地图")
        self.ax.axis('off')
        
        # 更新画布
        self.canvas.draw()
    
    def run_algorithm(self, algo_type):
        start = self.start_var.get()
        goal = self.goal_var.get()
        
        # 验证起点和终点
        if start not in self.romania_map:
            messagebox.showerror("错误", f"起点 '{start}' 不在地图中!")
            return
            
        if goal not in self.heuristic:
            messagebox.showerror("错误", f"终点 '{goal}' 不在启发函数中!")
            return
        
        # 运行选择的算法
        if algo_type == "ucs":
            path, cost, expanded = uniform_cost_search(self.romania_map, start, goal)
            algo_name = "代价一致搜索"
        elif algo_type == "greedy":
            path, cost, expanded = greedy_search(self.romania_map, start, goal, self.heuristic)
            algo_name = "贪婪最佳优先搜索"
        elif algo_type == "astar":
            path, cost, expanded = a_star_search(self.romania_map, start, goal, self.heuristic)
            algo_name = "A*算法"
        
        # 显示结果
        self.result_text.delete(1.0, tk.END)
        if path:
            self.result_text.insert(tk.END, f"{algo_name}结果:\n\n")
            self.result_text.insert(tk.END, f"路径: {' -> '.join(path)}\n\n")
            self.result_text.insert(tk.END, f"总代价: {cost}\n\n")
            self.result_text.insert(tk.END, f"扩展节点数: {expanded}")
            
            # 在地图上标记路径
            self.draw_map(path)
        else:
            self.result_text.insert(tk.END, f"{algo_name}未找到路径!\n\n")
            self.result_text.insert(tk.END, f"扩展节点数: {expanded}")
            
            # 重绘地图
            self.draw_map()

# 主程序入口
def main():
    root = tk.Tk()
    app = RomaniaMapApp(root)
    root.mainloop()

if __name__ == "__main__":
    main()